import sys
import threading
import cv2
from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtWidgets import (QLabel, QPushButton, QVBoxLayout, QHBoxLayout, 
                            QWidget, QApplication, QFrame, QMessageBox)
from PyQt5.QtGui import QPixmap, QImage, QFont, QColor
from PyQt5.QtCore import QTimer
from test_detect import FruitSortingSystem, FruitDetector, FruitGrabber

class FruitSorterGUI(QWidget):
    def __init__(self, bg_path):
        super().__init__()
        self.bg_path = bg_path
        self.system = None  # Will be initialized when "Start System" is clicked
        self.detector = None
        self.grabber = None
        self.init_ui()
        self.camera_on = False
        self.detecting = False
        self.grabbing = False
        self.obstacle_mode = False
        self.capture_timer = QTimer()
        self.capture_timer.timeout.connect(self.update_frame)
        self.frame = None

    def init_ui(self):
        self.setWindowTitle("Fruit Sorting Robot Arm")
        self.setFixedSize(1080, 720)

        # Set background image
        self.bg_label = QLabel(self)
        self.bg_label.setGeometry(0, 0, 1080, 720)
        pixmap = QPixmap(self.bg_path).scaled(self.size(), QtCore.Qt.KeepAspectRatioByExpanding)
        self.bg_label.setPixmap(pixmap)
        self.bg_label.lower()

        # Camera frame
        self.cam_frame = QLabel(self)
        self.cam_frame.setGeometry(120, 100, 600, 400)
        self.cam_frame.setStyleSheet("border: 4px solid #FF9800; background-color: #333333;")
        self.cam_frame.setAlignment(QtCore.Qt.AlignCenter)
        self.cam_frame.setText("Camera Feed")

        # Title
        self.title = QLabel("Fruit Sorting Helper", self)
        self.title.setGeometry(0, 20, 1080, 60)
        self.title.setFont(QFont("Berlin Sans FB Demi", 32, QFont.Bold))
        self.title.setStyleSheet("color: white; background: rgba(0,0,0,0.3);")
        self.title.setAlignment(QtCore.Qt.AlignCenter)

        # System control buttons (top row)
        self.start_system_btn = QPushButton("Initialize", self)
        self.start_system_btn.setGeometry(120, 530, 180, 60)
        self.start_system_btn.setFont(QFont("Arial", 18))
        self.start_system_btn.setStyleSheet("background: #2196F3; color: white; border-radius: 10px;")
        self.start_system_btn.clicked.connect(self.start_system)

        # Camera control buttons
        self.capture_btn = QPushButton("Set Mode", self)
        self.capture_btn.setGeometry(350, 530, 180, 60)
        self.capture_btn.setFont(QFont("Arial", 18))
        self.capture_btn.setStyleSheet("background: #2196F3; color: white; border-radius: 10px;")
        self.capture_btn.clicked.connect(self.capture_frame)
        self.capture_btn.setEnabled(False)

        self.detect_btn = QPushButton("Start Detection", self)
        self.detect_btn.setGeometry(580, 530, 180, 60)
        self.detect_btn.setFont(QFont("Arial", 18))
        self.detect_btn.setStyleSheet("background: #4CAF50; color: white; border-radius: 10px;")
        self.detect_btn.clicked.connect(self.toggle_detection)
        self.detect_btn.setEnabled(False)


        # Detection and obstacle buttons
        self.camera_btn = QPushButton("Open Camera", self)
        self.camera_btn.setGeometry(120, 610, 180, 60)
        self.camera_btn.setFont(QFont("Arial", 18))
        self.camera_btn.setStyleSheet("background: #4CAF50; color: white; border-radius: 10px;")
        self.camera_btn.clicked.connect(self.toggle_camera)
        self.camera_btn.setEnabled(False)

        self.obstacle_btn = QPushButton("Set Obstacle", self)
        self.obstacle_btn.setGeometry(350, 610, 180, 60)
        self.obstacle_btn.setFont(QFont("Arial", 18))
        self.obstacle_btn.setStyleSheet("background: #FF9800; color: white; border-radius: 10px;")
        self.obstacle_btn.clicked.connect(self.set_obstacle)
        self.obstacle_btn.setEnabled(False)

        # Grab control buttons
        self.grab_btn = QPushButton("Start Grab", self)
        self.grab_btn.setGeometry(580, 610, 180, 60)
        self.grab_btn.setFont(QFont("Arial", 18))
        self.grab_btn.setStyleSheet("background: #9C27B0; color: white; border-radius: 10px;")
        self.grab_btn.clicked.connect(self.start_grab)
        self.grab_btn.setEnabled(False)

        self.exit_btn = QPushButton("Exit", self)
        self.exit_btn.setGeometry(800, 610, 180, 60)
        self.exit_btn.setFont(QFont("Arial", 18))
        self.exit_btn.setStyleSheet("background: #607D8B; color: white; border-radius: 10px;")
        self.exit_btn.clicked.connect(self.close)

        # Fruit counts panel
        self.count_panel = QLabel(self)
        self.count_panel.setGeometry(760, 120, 250, 300)
        self.count_panel.setFont(QFont("Consolas", 16))
        self.count_panel.setAlignment(QtCore.Qt.AlignTop)
        self.count_panel.setStyleSheet("background: rgba(255,255,255,0.75); border-radius: 18px; padding: 12px;")

        # Status bar
        self.status_label = QLabel("Ready", self)
        self.status_label.setGeometry(0, 680, 1080, 36)
        self.status_label.setFont(QFont("Arial", 16))
        self.status_label.setStyleSheet("background: #37474F; color: #FFF;")
        self.status_label.setAlignment(QtCore.Qt.AlignCenter)

        self.update_counts([0]*8)

    # Button functions implementation
    def start_system(self):
        """1. 启动系统：等同于test_detect中的system=FruitSortingSystem()"""
        try:
            self.system = FruitSortingSystem()
            self.detector = self.system.detector
            self.grabber = self.system.grabber
            self.status_label.setText("Initializing...")
            QApplication.processEvents()  # 强制立即处理UI更新
            if self.system.start():
                self.status_label.setText("Initialization completed successfully")
            else:
                self.status_label.setText("Initialization process encountered an error")
            self.camera_btn.setEnabled(True)            
            self.start_system_btn.setEnabled(False)
        except Exception as e:
            self.status_label.setText(f"Failed to start system: {str(e)}")
            QMessageBox.critical(self, "Error", f"Failed to start system: {str(e)}")

    def toggle_camera(self):
        """2. 打开摄像头：打开摄像头，提示放置卡片，并激活拍摄按钮"""
        if not self.camera_on:
            if self.detector.open_camera():
                self.camera_on = True
                self.camera_btn.setText("Close Camera")
                self.capture_btn.setEnabled(True)
                #self.close_camera_btn.setEnabled(True)
                self.detect_btn.setEnabled(True)
                self.status_label.setText("Camera opened - Please place the card")
                
                # Start camera preview
                self.capture_timer.start(30)  # Update every 30ms
            else:
                self.status_label.setText("Failed to open camera")
        else:
            self.close_camera()

    def capture_frame(self):
        """3. 拍摄：获取摄像头帧，并询问是否设置障碍物"""
        if not self.camera_on:
            return
            
        self.frame = self.detector.get_frame()
        frame = self.frame
        if frame is not None:
            # Show the captured frame
            self._update_frame(frame)
            
            # Ask about obstacle mode
            reply = QMessageBox.question(self, 'Obstacle Mode', 
                                       'Do you want to set obstacle mode?', 
                                       QMessageBox.Yes | QMessageBox.No, 
                                       QMessageBox.No)
            
            if reply == QMessageBox.Yes:
                self.obstacle_mode = True
                self.detector.set_obstacle_mode(True)
                self.obstacle_btn.setEnabled(True)
                self.status_label.setText("Obstacle mode activated - Please set obstacle point")
            else:
                self.obstacle_mode = False
                self.detector.set_obstacle_mode(False)
                self.obstacle_btn.setEnabled(False)
                self.status_label.setText("Normal mode activated")

    def close_camera(self):
        """4. 关闭摄像头"""
        if self.camera_on:
            self.capture_timer.stop()
            self.detector.close_camera()
            self.frame = None
            self.camera_on = False
            self.detecting = False
            self.camera_btn.setText("Open Camera")
            self.capture_btn.setEnabled(False)
            # self.close_camera_btn.setEnabled(False)
            self.detect_btn.setEnabled(False)
            self.obstacle_btn.setEnabled(False)
            self.grab_btn.setEnabled(False)
            self.status_label.setText("Camera closed")
            self.cam_frame.setText("Camera Feed")

    def set_obstacle(self):
        """5. 设置障碍物：矩形检测设置障碍点"""
        if not self.camera_on:
            return
            
        frame = self.detector.get_frame()
        if frame is not None:
            # Use the rectangle detector to find obstacle point
            if self.system.set_obstacle_from_rectangle(frame):
                obstacle_point = self.detector.obstacle_point
                if obstacle_point:
                    # Draw the obstacle point on the frame
                    cv2.circle(frame, (obstacle_point[0], obstacle_point[1]), 10, (0, 0, 255), -1)
                    self._update_frame(frame)
                    self.status_label.setText(f"Obstacle set at {obstacle_point}")
                    
                    # Enable grab button if not already enabled
                    if not self.grab_btn.isEnabled():
                        self.grab_btn.setEnabled(True)
            else:
                self.status_label.setText("Failed to detect obstacle point")

    def toggle_detection(self):
        """6. 开始检测：若检测到水果，在相机框出绘制出检测结果"""
        if not self.detecting:
            if self.detector.start_detection():
                self.detecting = True
                self.detect_btn.setText("Stop Detection")
                self.status_label.setText("Detection started")
                
                # Disable some buttons during detection
                self.capture_btn.setEnabled(False)
                self.obstacle_btn.setEnabled(False)
                
                # Start detection loop
                self.capture_timer.timeout.disconnect()
                self.capture_timer.timeout.connect(self.detection_loop)
            else:
                self.status_label.setText("Failed to start detection")
        else:
            self.stop_detection()

    def detection_loop(self):
        """Detection processing loop"""
        if not self.detecting or not self.camera_on:
            return
            
        frame = self.detector.get_frame()
        if frame is not None:
            # Perform detection
            boxes, result_img = self.detector.detect_fruits(frame)
            
            # Update the display
            self._update_frame(result_img)
            self.update_counts(self.detector.cnt)
            
            # Check if we should grab
            grab_target = self.detector.should_grab(boxes)
            if grab_target:
                self.status_label.setText("Fruit detected - Ready to grab")
                self.grab_btn.setEnabled(True)

    def start_grab(self):
        """7. 抓取按钮：执行抓取动作"""
        if not self.detecting or not self.camera_on:
            return
        
        frame = self.detector.get_frame()
        if frame is None:
            return
        
        self.capture_timer.stop()  # 暂停画面更新
            
        # Get the latest detection results
        boxes, _ = self.detector.detect_fruits(frame)
        grab_target = self.detector.should_grab(boxes)
        
        if grab_target:
            label = int(grab_target[5])
            try:
                self.grabbing = True
                self.grab_btn.setEnabled(False)
                self.status_label.setText("Grabbing in progress...")
                
                # Perform grab in a separate thread to avoid freezing GUI
                grab_thread = threading.Thread(
                    target=self.grabber.grab_fruit,
                    args=(grab_target, label, self.obstacle_mode)
                )
                grab_thread.start()
                
                # Wait for grab to complete (in a non-blocking way)
                self.check_grab_status(grab_thread)
                self.detector.stop_detection()
            except Exception as e:
                self.status_label.setText(f"Grab failed: {str(e)}")
                self.detector.stop_detection()
                self.grabbing = False
                self.grab_btn.setEnabled(True)
        else:
            self.status_label.setText("No valid target to grab")

    def check_grab_status(self, thread):
        """Check if grab thread has finished"""
        if thread.is_alive():
            QTimer.singleShot(100, lambda: self.check_grab_status(thread))
        else:
            self.grabbing = False
            self.status_label.setText("Grab completed")
            self.capture_timer.start(30)  # 恢复画面更新
            
            # Re-enable detection if it was running
            if self.detecting:
                self.grab_btn.setEnabled(True)

    def stop_detection(self):
        """Stop the detection process"""
        self.detector.stop_detection()
        self.detecting = False
        self.detect_btn.setText("Start Detection")
        self.status_label.setText("Detection stopped")
        
        # Re-enable buttons
        self.capture_btn.setEnabled(True)
        if self.obstacle_mode:
            self.obstacle_btn.setEnabled(True)
        
        # Switch back to camera preview
        self.capture_timer.timeout.disconnect()
        self.capture_timer.timeout.connect(self.update_frame)

    def update_frame(self):
        """Update camera preview frame"""
        if not self.camera_on:
            return
            
        frame = self.detector.get_frame()
        if frame is not None:
            self._update_frame(frame)

    def _update_frame(self, frame):
        """Helper to update frame display"""
        try:
            rgb_image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            h, w, ch = rgb_image.shape
            bytes_per_line = ch * w
            qt_image = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format_RGB888)
            pixmap = QPixmap.fromImage(qt_image).scaled(
                self.cam_frame.width(), self.cam_frame.height(), 
                QtCore.Qt.KeepAspectRatio)
            self.cam_frame.setPixmap(pixmap)
        except Exception as e:
            print(f"Frame update error: {str(e)}")

    def update_counts(self, counts):
        fruit_names = [
            "Defect", "Fresh Apple", "Fresh Banana", "Fresh Orange",
            "Mold", "Rotten Apple", "Rotten Banana", "Rotten Orange"
        ]
        html = "<b>Fruit Counts:</b><br>"
        for i, name in enumerate(fruit_names):
            html += f"{name}: <b style='color:#4CAF50'>{counts[i]}</b><br>"
        self.count_panel.setText(html)
    
    def closeEvent(self, event):
        """Clean up resources when closing window"""
        if self.camera_on:
            self.close_camera()
        
        if self.system:
            try:
                self.status_label.setText("Release resource")
                QApplication.processEvents()
                self.system.stop()
            except Exception as e:
                print(f"Cleanup error: {str(e)}")
        
        event.accept()



if __name__ == "__main__":
    app = QApplication(sys.argv)
    bg_path = "bg.jpg"  # Set your background image path
    window = FruitSorterGUI(bg_path)
    window.show()
    sys.exit(app.exec_())